# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
from typing import List
from mesh import MeshBlock
from Loads import loadBase,BcBase,ConcentratedLoad,BodyLoad, PressureLoad, SurfaceLoad, DisplacementBc
from tools import *
# 必须只有一种单元格式
class FEA:
    """将MeshBlock, BCs, loads等信息整合到一起,进行求解计算"""
    def __init__(self, meshblocks: List[MeshBlock], 
                bcs: List[BcBase], 
                loads: List[loadBase]):
        """初始化FEA类, 传入MeshBlock, BCs, loads"""
        self.GlobalMeshblock:MeshBlock = sum(meshblocks)
        self.BcList = bcs
        self.LoadList = loads
        # 汇总各个MeshBlock的自由度个数
        self.GlobalDofNum = self.GlobalMeshblock.BlockDofNum
        # VTK elem type
        self.VtkType=self.GlobalMeshblock.VTK_TYPE
        self.initResultVar()
    def initResultVar(self)->None:
        """初始化变量U,Fg等"""
        self._Ug=np.zeros((self.GlobalDofNum,1))
        self._F_react=np.zeros((self.GlobalDofNum,1))
        self._ConstraingedDofIndexs:List[int]=[]
    @property
    def Ug(self)->np.ndarray:
        """返回全局位移向量(N,1)"""
        return self._Ug
    @Ug.setter
    def Ug(self,value:np.ndarray)->None:
        """设置全局位移向量"""
        if value.shape!=(self.GlobalDofNum,1):
            raise ValueError("Ug shape is not correct")
        self._Ug=value
    @property
    def F_react(self)->np.ndarray:
        """返回全局恢复力(N,1)"""
        return self._F_react
    @F_react.setter
    def F_react(self,value:np.ndarray)->None:
        """设置全局恢复力"""
        if value.shape!=(self.GlobalDofNum,1):
            raise ValueError("F_react shape is not correct")
        self._F_react=value
    def __str__(self)->str:
        """打印FEA信息"""
        return f"FEA:\nglobalDofNum={self.GlobalDofNum}"
    @property
    def Kg(self)->np.ndarray:
        """系统全局刚度矩阵"""
        return self.GlobalMeshblock.Kg(globalDofNum=self.GlobalDofNum)
    def showKg(self)->None:
        """imshow显示全局刚度矩阵"""
        imShowMat(self.Kg,"global stiffness matrix")
    def showMg(self)->None:
        """imshow显示全局质量矩阵"""
        imShowMat(self.Mg,"global mass matrix")
    @property
    def Mg(self)->np.ndarray:   
        """系统全局质量矩阵"""
        return self.GlobalMeshblock.Mg(globalDofNum=self.GlobalDofNum)
    def PreProcessFg(self)->Tuple[np.ndarray,Dict[int,float]]:
        """将各个载荷转换为等效节点力,汇总为全局外力向量,返回值为全局加载节点力向量(n,1); 
        全局节点力字典{key:自由度索引(0-base), value:节点力}
        """
        Fg=np.zeros((self.GlobalDofNum,1))
        allElems=self.GlobalMeshblock.ElemDict
        allnodes=self.GlobalMeshblock.NodesDict
        for load in self.LoadList:
            if load.isElementLoad:
                # 处理作用于单元的载荷,循环所有单元,如果不在load则continue
                for elabel in allElems.keys():
                    if elabel not in load.get_ElemsList:
                        continue
                    # 将单元等效节点力加到单元的自由度索引位置的全局外力向量Fg中,
                    # print(Fg[allElems[elabel].ElemDofIndexs,0].shape)
                    Fg[allElems[elabel].ElemDofIndexs,0]+=allElems[elabel].ProcessELoad2EqNodeForce(load=load)[:,0]
                print(f"-->{load.__class__.__name__} 载荷 {load.name}已处理")
            else:
                # 处理作用于节点的载荷
                if isinstance(load,ConcentratedLoad):
                    for nlabel in allnodes.keys():
                        if nlabel not in load.get_NodesList:
                            continue
                        ind=get_ithdof_index(NodeLabel=nlabel,
                                            dofNum=self.GlobalMeshblock.DofPerNode,
                                            ithdof=load.dofth)
                        Fg[ind,0]+=load.LoadValue
                    print(f"-->{load.__class__.__name__} 载荷 {load.name}已处理")
        
        # 找出Fg中非0值的索引和非零值,定义全局节点力字典{key:自由度索引(0-base), value:节点力}
        P_ind=Fg.nonzero()
        P_dict=dict(zip(P_ind[0].tolist(),Fg[P_ind].tolist()))
        return Fg,P_dict
    def ProcessBcs(self)->Tuple[np.ndarray,np.ndarray,np.ndarray,Dict[int,float]]:
        """将各个边界条件转换为等效节点位移,汇总为全局位移向量;
        
        返回Kg_,Mg_,Fg_,BcDic,其中:前三个是已经用充大数法处理过的Kg,Mg,Fg;BcDic是全局边界条件字典{key:约束的自由度索引; value:位移约束值}
        """
        # 计算全局刚度,质量矩阵,节点力向量
        Kg_=self.Kg.copy()
        Mg_=self.Mg.copy()
        Fg_=self.PreProcessFg()[0].copy()
        # 处理边界条件(充大数法)
        BcDic:Dict[int,float]={}  #key:约束的自由度索引; value:约束值
        allnodes=self.GlobalMeshblock.NodesDict
        for bc in self.BcList:
            if isinstance(bc,DisplacementBc):
                # 强制处理位移边界条件
                for nlabel in  allnodes.keys():
                    if nlabel not in bc.get_NodesList:
                        continue
                    ind=get_ithdof_index(NodeLabel=nlabel,
                                        dofNum=self.GlobalMeshblock.DofPerNode,
                                        ithdof=bc.Dofth)
                    if ind not in BcDic.keys():
                        BcDic[ind]=bc.Value
                    else:
                        BcDic[ind]+=bc.Value
            print(f"-->{bc.__class__.__name__} 边界条件 {bc.name}已处理")
        # 使用充大数法处理Kg_,Mg_,Fg_
        C=np.max(Kg_+Mg_)*10e6
        for ind,val in BcDic.items():
            Kg_[ind,ind]+=C
            Mg_[ind,ind]+=C
            Fg_[ind,0]+=val*C
        self._ConstraingedDofIndexs=list(BcDic.keys())
        return Kg_,Mg_,Fg_,BcDic
    def Solve(self,**kwargs):
        """求解系统方程,并更新Ug,Fg等结果
        input:
            StepType: 分析类型,目前只支持静态分析
        """
        if kwargs['StepType'] == 'Static':
            Kg_,Mg_,Fg_,_=self.ProcessBcs()
            # 静态分析,解出一个Ug的列向量
            self.Ug=np.linalg.solve(Kg_,Fg_)
            self.F_react=self.Kg@self.Ug
        return self.Ug,self.F_react
